#pragma comment(lib, "Kinect10")

#include <Windows.h>
#include <NuiApi.h>

HANDLE ColorStreamHandle = INVALID_HANDLE_VALUE;
HANDLE DepthStreamHandle = INVALID_HANDLE_VALUE;

const NUI_IMAGE_FRAME * CurrentColorFrame;
const NUI_IMAGE_FRAME * CurrentDepthFrame;

NUI_LOCKED_RECT CurrentColorRect;
NUI_LOCKED_RECT CurrentDepthRect;

DWORD CurrentColorBuffer[1280 * 960];
WORD  CurrentDepthBuffer[640 * 480];

int Resolution[2] = {640, 480};


bool KinectInput_getCurrentFrame(DWORD * * colorFrame, WORD * * depthFrame, bool mirror) {
	if (ColorStreamHandle == INVALID_HANDLE_VALUE) {
		return false;
	}
	if (NuiImageStreamGetNextFrame(ColorStreamHandle, 50, &CurrentColorFrame) != S_OK) {
		return false;
	}
	if (DepthStreamHandle != INVALID_HANDLE_VALUE && NuiImageStreamGetNextFrame(DepthStreamHandle, 50, &CurrentDepthFrame) != S_OK) {
		return false;
	}
	if (CurrentColorFrame->pFrameTexture->LockRect(0, &CurrentColorRect, 0, 0) != S_OK) {
		return false;
	}
	if (DepthStreamHandle != INVALID_HANDLE_VALUE && CurrentDepthFrame->pFrameTexture->LockRect(0, &CurrentDepthRect, 0, 0) != S_OK) {
		return false;
	}
	{
		int lines = CurrentColorRect.size / CurrentColorRect.Pitch;
		int rows = CurrentColorRect.Pitch / 4;
		unsigned long * tmpRect = (unsigned long *)CurrentColorRect.pBits;
		for (int i = 0; i < lines; i++) {
			for (int j = 0; j < rows; j++) {
				if (mirror) {
					CurrentColorBuffer[i * rows + j] = tmpRect[i * rows + j];
				} else {
					CurrentColorBuffer[i * rows + j] = tmpRect[i * rows + (rows - 1 - j)];
				}
			}
		}
	}
	if (colorFrame != NULL) {
		*colorFrame = CurrentColorBuffer;
	}
	
	memcpy(CurrentDepthBuffer, CurrentDepthRect.pBits, CurrentDepthRect.size);
	if (DepthStreamHandle != INVALID_HANDLE_VALUE && depthFrame != NULL) {
		*depthFrame = CurrentDepthBuffer;
	}
	CurrentColorFrame->pFrameTexture->UnlockRect(0);
	NuiImageStreamReleaseFrame(ColorStreamHandle, CurrentColorFrame);
	if (DepthStreamHandle != INVALID_HANDLE_VALUE) {
		CurrentDepthFrame->pFrameTexture->UnlockRect(0);
		NuiImageStreamReleaseFrame(DepthStreamHandle, CurrentDepthFrame);
	}
	return true;
}

bool KinectInput_init(bool highRes, bool depth) {
	HRESULT ret;
	NUI_IMAGE_RESOLUTION resType = NUI_IMAGE_RESOLUTION_640x480;

	if (highRes) {
		resType = NUI_IMAGE_RESOLUTION_1280x960;
		Resolution[0] = 1280;
		Resolution[1] = 720;
	}
	if (depth) {
		if (ret = NuiInitialize(NUI_INITIALIZE_FLAG_USES_DEPTH | NUI_INITIALIZE_FLAG_USES_COLOR) != S_OK) { 
			return false;
		}
	} else {
		ret = NuiInitialize(NUI_INITIALIZE_FLAG_USES_COLOR);
		if (ret != S_OK) { 
			return false;
		}
	}
	ret = NuiImageStreamOpen(
		NUI_IMAGE_TYPE_COLOR,
		resType,
		NUI_IMAGE_STREAM_FLAG_SUPPRESS_NO_FRAME_DATA,
		2,
		0,
		&ColorStreamHandle);
	if (ret != S_OK)
	{
		return false;
	}
	if (depth) {
			ret = NuiImageStreamOpen(
			NUI_IMAGE_TYPE_DEPTH, 
			NUI_IMAGE_RESOLUTION_640x480, 
			NUI_IMAGE_STREAM_FLAG_SUPPRESS_NO_FRAME_DATA,
			2, 
			0,
			&DepthStreamHandle);
			if (ret != S_OK)
		{
			return false;
		}
	}
	if (ret != S_OK) {
		return false;
	}
	return true;
}

void KinectInput_finalize() {
	NuiShutdown();
}

void KinectInput_getResolution(int * width, int * height) {
	*width = Resolution[0];
	*height = Resolution[1];
}